from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Iterable, List, Set, Tuple

from src.db import query_pdf_ids_by_tag


Token = Tuple[str, str]  # (type, value) types: TAG, AND, OR, LPAREN, RPAREN


def tokenize(expr: str) -> List[Token]:
    s = expr.strip()
    if not s:
        return []
    tokens: List[Token] = []
    i = 0
    while i < len(s):
        ch = s[i]
        if ch.isspace():
            i += 1
            continue
        if ch == '(':  # lparen
            tokens.append(("LPAREN", ch))
            i += 1
            continue
        if ch == ')':  # rparen
            tokens.append(("RPAREN", ch))
            i += 1
            continue
        if ch in ('"', "'"):
            quote = ch
            i += 1
            start = i
            while i < len(s) and s[i] != quote:
                i += 1
            val = s[start:i]
            i += 1 if i < len(s) else 0
            tokens.append(("TAG", val))
            continue
        # word
        start = i
        while i < len(s) and not s[i].isspace() and s[i] not in '()':
            i += 1
        word = s[start:i]
        upper = word.upper()
        if upper == 'AND':
            tokens.append(("AND", word))
        elif upper == 'OR':
            tokens.append(("OR", word))
        else:
            tokens.append(("TAG", word))
    return tokens


# Shunting-yard to RPN
def _validate(tokens: List[Token]) -> None:
    depth = 0
    expect_operand = True
    for ttype, _ in tokens:
        if ttype == "TAG":
            if not expect_operand:
                raise ValueError("Unexpected tag")
            expect_operand = False
        elif ttype == "LPAREN":
            if not expect_operand:
                raise ValueError("Unexpected '('")
            depth += 1
        elif ttype == "RPAREN":
            if expect_operand:
                raise ValueError("Unexpected ')'")
            depth -= 1
            if depth < 0:
                raise ValueError("Unmatched ')' in expression")
        elif ttype in ("AND", "OR"):
            if expect_operand:
                raise ValueError("Operator without preceding operand")
            expect_operand = True
    if depth != 0:
        raise ValueError("Unmatched '(' in expression")
    if expect_operand:
        raise ValueError("Expression ended unexpectedly")


def to_rpn(tokens: List[Token]) -> List[Token]:
    prec = {"AND": 2, "OR": 1}
    output: List[Token] = []
    stack: List[Token] = []
    for ttype, val in tokens:
        if ttype == "TAG":
            output.append((ttype, val))
        elif ttype in ("AND", "OR"):
            while stack and stack[-1][0] in prec and prec[stack[-1][0]] >= prec[ttype]:
                output.append(stack.pop())
            stack.append((ttype, val))
        elif ttype == "LPAREN":
            stack.append((ttype, val))
        elif ttype == "RPAREN":
            while stack and stack[-1][0] != "LPAREN":
                output.append(stack.pop())
            if stack and stack[-1][0] == "LPAREN":
                stack.pop()
    while stack:
        output.append(stack.pop())
    return output


def eval_rpn(conn, rpn: List[Token]) -> Set[int]:
    stack: List[Set[int]] = []
    for ttype, val in rpn:
        if ttype == "TAG":
            ids = query_pdf_ids_by_tag(conn, val)
            stack.append(ids)
        elif ttype == "AND":
            b = stack.pop() if stack else set()
            a = stack.pop() if stack else set()
            stack.append(a & b)
        elif ttype == "OR":
            b = stack.pop() if stack else set()
            a = stack.pop() if stack else set()
            stack.append(a | b)
    return stack[-1] if stack else set()


def search_pdf_ids(conn, expr: str) -> Set[int]:
    tokens = tokenize(expr)
    if not tokens:
        return set()
    _validate(tokens)
    rpn = to_rpn(tokens)
    return eval_rpn(conn, rpn)


